package com.lz.tree.rb;

public class RedBlackTree {
    private RBNode root;
    private final RBNode NIL; // 哨兵节点，代表叶子节点（NIL）

    // 构造函数
    public RedBlackTree() {
        NIL = new RBNode(-1); // 哨兵节点键值任意，不会被使用
        NIL.color = 0;        // 叶子节点必须是黑色
        root = NIL;           // 初始化根节点为哨兵
    }

    // 左旋操作
    private void leftRotate(RBNode x) {
        RBNode y = x.right;   // 设置y为x的右子节点
        x.right = y.left;     // 将y的左子树变为x的右子树

        if (y.left != NIL) {
            y.left.parent = x; // 更新y左子节点的父指针
        }

        y.parent = x.parent;  // 将x的父节点赋给y

        if (x.parent == NIL) {
            root = y;         // 如果x是根节点，则y成为新根
        } else if (x == x.parent.left) {
            x.parent.left = y; // 如果x是左子节点，则y成为左子节点
        } else {
            x.parent.right = y; // 如果x是右子节点，则y成为右子节点
        }

        y.left = x;           // 将x放在y的左边
        x.parent = y;         // 更新x的父节点
    }

    // 右旋操作（对称于左旋）
    private void rightRotate(RBNode y) {
        RBNode x = y.left;
        y.left = x.right;

        if (x.right != NIL) {
            x.right.parent = y;
        }

        x.parent = y.parent;

        if (y.parent == NIL) {
            root = x;
        } else if (y == y.parent.right) {
            y.parent.right = x;
        } else {
            y.parent.left = x;
        }

        x.right = y;
        y.parent = x;
    }

    // 插入节点
    public void insert(int key) {
        RBNode z = new RBNode(key);
        RBNode y = NIL;
        RBNode x = root;

        // 查找插入位置
        while (x != NIL) {
            y = x;
            if (z.key < x.key) {
                x = x.left;
            } else {
                x = x.right;
            }
        }

        z.parent = y;
        if (y == NIL) {
            root = z;         // 树为空时，新节点为根
        } else if (z.key < y.key) {
            y.left = z;
        } else {
            y.right = z;
        }

        // 初始化新节点的子节点为哨兵
        z.left = NIL;
        z.right = NIL;
        z.color = 1;          // 新节点初始为红色

        insertFixup(z);       // 修复插入可能破坏的红黑树性质
    }

    // 插入后修复红黑树性质
    private void insertFixup(RBNode z) {
        // 当z的父节点是红色时需要修复（违反性质4）
        while (z.parent.color == 1) {
            if (z.parent == z.parent.parent.left) {
                RBNode y = z.parent.parent.right; // 叔节点

                // 情况1：叔节点为红色
                if (y.color == 1) {
                    z.parent.color = 0;         // 父节点变黑
                    y.color = 0;                // 叔节点变黑
                    z.parent.parent.color = 1;   // 祖父节点变红
                    z = z.parent.parent;        // z上移到祖父节点
                } else {
                    // 情况2：z是右子节点（三角关系）
                    if (z == z.parent.right) {
                        z = z.parent;           // z上移
                        leftRotate(z);          // 左旋转为直线关系
                    }
                    // 情况3：z是左子节点（直线关系）
                    z.parent.color = 0;           // 父节点变黑
                    z.parent.parent.color = 1;    // 祖父节点变红
                    rightRotate(z.parent.parent); // 右旋祖父节点
                }
            } else { // 对称情况（父节点是右子节点）
                RBNode y = z.parent.parent.left; // 叔节点
                if (y.color == 1) {
                    z.parent.color = 0;
                    y.color = 0;
                    z.parent.parent.color = 1;
                    z = z.parent.parent;
                } else {
                    if (z == z.parent.left) {
                        z = z.parent;
                        rightRotate(z);
                    }
                    z.parent.color = 0;
                    z.parent.parent.color = 1;
                    leftRotate(z.parent.parent);
                }
            }
        }
        root.color = 0; // 根节点始终为黑
    }

    // 查找最小节点
    private RBNode minimum(RBNode node) {
        while (node.left != NIL) {
            node = node.left;
        }
        return node;
    }

    // 节点替换（用v替换u在树中的位置）
    private void transplant(RBNode u, RBNode v) {
        if (u.parent == NIL) {
            root = v;
        } else if (u == u.parent.left) {
            u.parent.left = v;
        } else {
            u.parent.right = v;
        }
        v.parent = u.parent;
    }

    // 删除节点
    public void delete(int key) {
        RBNode z = search(key);
        if (z == NIL) return; // 节点不存在

        RBNode y = z;
        RBNode x;
        int yOriginalColor = y.color;

        if (z.left == NIL) {
            x = z.right;
            transplant(z, z.right);
        } else if (z.right == NIL) {
            x = z.left;
            transplant(z, z.left);
        } else {
            y = minimum(z.right); // 找到后继节点
            yOriginalColor = y.color;
            x = y.right;

            if (y.parent == z) {
                x.parent = y;
            } else {
                transplant(y, y.right);
                y.right = z.right;
                y.right.parent = y;
            }

            transplant(z, y);
            y.left = z.left;
            y.left.parent = y;
            y.color = z.color; // 继承被删除节点的颜色
        }

        // 如果被删除的原始节点是黑色，需要修复
        if (yOriginalColor == 0) {
            deleteFixup(x);
        }
    }

    // 删除后修复红黑树性质
    private void deleteFixup(RBNode x) {
        while (x != root && x.color == 0) {
            if (x == x.parent.left) {
                RBNode w = x.parent.right; // 兄弟节点

                // 情况1：兄弟节点是红色
                if (w.color == 1) {
                    w.color = 0;          // 兄弟节点变黑
                    x.parent.color = 1;    // 父节点变红
                    leftRotate(x.parent);  // 左旋父节点
                    w = x.parent.right;    // 更新兄弟节点
                }

                // 情况2：兄弟节点的两个子节点都是黑色
                if (w.left.color == 0 && w.right.color == 0) {
                    w.color = 1;          // 兄弟节点变红
                    x = x.parent;         // 向上传递问题
                } else {
                    // 情况3：兄弟节点的左子节点红，右子节点黑
                    if (w.right.color == 0) {
                        w.left.color = 0; // 兄弟左子节点变黑
                        w.color = 1;       // 兄弟节点变红
                        rightRotate(w);    // 右旋兄弟节点
                        w = x.parent.right; // 更新兄弟节点
                    }

                    // 情况4：兄弟节点的右子节点红
                    w.color = x.parent.color; // 兄弟节点继承父节点颜色
                    x.parent.color = 0;       // 父节点变黑
                    w.right.color = 0;       // 兄弟右子节点变黑
                    leftRotate(x.parent);     // 左旋父节点
                    x = root;                 // 结束循环
                }
            } else { // 对称情况（x是右子节点）
                RBNode w = x.parent.left;
                if (w.color == 1) {
                    w.color = 0;
                    x.parent.color = 1;
                    rightRotate(x.parent);
                    w = x.parent.left;
                }
                if (w.right.color == 0 && w.left.color == 0) {
                    w.color = 1;
                    x = x.parent;
                } else {
                    if (w.left.color == 0) {
                        w.right.color = 0;
                        w.color = 1;
                        leftRotate(w);
                        w = x.parent.left;
                    }
                    w.color = x.parent.color;
                    x.parent.color = 0;
                    w.left.color = 0;
                    rightRotate(x.parent);
                    x = root;
                }
            }
        }
        x.color = 0; // 最终确保x为黑色
    }

    // 查找节点
    private RBNode search(int key) {
        RBNode current = root;
        while (current != NIL && key != current.key) {
            if (key < current.key) {
                current = current.left;
            } else {
                current = current.right;
            }
        }
        return current;
    }

    // 中序遍历（测试用）
    public void inorder() {
        inorderRec(root);
    }

    private void inorderRec(RBNode node) {
        if (node != NIL) {
            inorderRec(node.left);
            System.out.print(node.key + (node.color == 0 ? "B " : "R "));
            inorderRec(node.right);
        }
    }

    // 测试代码
    public static void main(String[] args) {
        RedBlackTree tree = new RedBlackTree();

        // 插入测试
        tree.insert(10);
        tree.insert(20);
        tree.insert(5);
        tree.insert(15);
        tree.insert(25);
        tree.insert(3);
        tree.insert(7);

        System.out.println("插入后中序遍历:");
        //tree.inorder(); // 输出: 3R 5B 7R 10B 15R 20B 25R

        tree.insert(16);
        tree.insert(17);
        // 删除测试
//        tree.delete(15);
//        System.out.println("\n删除15后中序遍历:");
//        tree.inorder(); // 输出: 3R 5B 7R 10B 20R 25B
//
//        tree.delete(5);
//        System.out.println("\n删除5后中序遍历:");
//        tree.inorder(); // 输出: 3B 7R 10B 20R 25B
    }
}
